import argparse
import os
import random
import sys
import time

import dgl
import numpy as np
import torch
import torch.nn.functional as F

from model import HINormer
from params import dblp_params
from utils.data import load_data
from utils.tools import mat2tensor, set_features_by_type
from utils.pytorchtools import EarlyStopping

# 这行代码用于设置环境变量CUDA_LAUNCH_BLOCKING为1。它强制CUDA代码同步执行，即每一个CUDA操作都会在控制台立即执行完毕并返回，而不是异步执行。这对于调试GPU相关代码非常有用，可以帮助找出某些异步错误或未及时执行的CUDA操作。
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

# 这行代码将'utils/'路径添加到Python的模块搜索路径中。这样可以从utils/目录中导入模块和文件，像你在代码中导入utils.data和utils.pytorchtools那样。
sys.path.append('utils/')


def run_model(args):
    # 检查是否存在checkpoint/目录，如果不存在则创建该目录。这个目录通常用于保存模型的中间检查点或结果文件。
    if not os.path.exists('checkpoint/'):
        os.makedirs('checkpoint/')

    # 调用load_data()函数加载数据集，返回：
    # features_list: 一个包含不同节点类型特征矩阵的列表。
    # adjM: 图的邻接矩阵。
    # labels: 节点的标签。
    # train_val_test_idx: 训练、验证和测试集的节点索引。
    # dl: 数据集加载器，用于处理训练/验证/测试过程中的评估和数据加载。
    features_list, adjM, labels, train_val_test_idx, dl = load_data(args.dataset)
    # 设置设备为GPU或CPU。如果GPU可用（torch.cuda.is_available()返回True），则选择用户指定的GPU设备cuda:device；否则，使用CPU。
    device = torch.device('cuda:' + str(args.device)
                          if torch.cuda.is_available() else 'cpu')
    # 遍历features_list中的每个特征矩阵，将其转换为PyTorch张量（通过mat2tensor()函数），并将其放置在指定的设备（CPU或GPU）上。
    features_list = [mat2tensor(features).to(device)
                     for features in features_list]
    # 计算每种节点类型的数量，features.shape[0]表示每个特征矩阵中的节点数量。
    node_cnt = [features.shape[0] for features in features_list]
    # 计算总节点数量，将所有节点的数量累加到sum_node。
    sum_node = 0
    for x in node_cnt:
        sum_node += x

    # 根据feats_type设置features_list
    # 因为不确定features_list会不会在函数里面被修改（值传递还是引用传递），因此这里多返回了一下
    in_dims, features_list = set_features_by_type(args.feats_type, features_list, device)

    # 将labels（节点的标签）转换为torch.LongTensor类型，并将其传送到指定的设备（GPU或CPU，取决于device）。
    labels = torch.LongTensor(labels).to(device)
    # 分别获取训练集、验证集、测试集的节点索引，并使用np.sort()对这些索引进行升序排序。
    train_idx = train_val_test_idx['train_idx']
    train_idx = np.sort(train_idx)
    val_idx = train_val_test_idx['val_idx']
    val_idx = np.sort(val_idx)
    test_idx = train_val_test_idx['test_idx']
    test_idx = np.sort(test_idx)

    # 创建一个DGL图对象g，邻接矩阵是adjM + adjM.T，表示将原始图的邻接矩阵和其转置相加。这样可以确保图是无向图（如果有某些边是单向的，相加会让它们变成双向的）。
    g = dgl.DGLGraph(adjM + (adjM.T))
    # 移除图中所有节点的自环（即从节点指向自身的边）。
    g = dgl.remove_self_loop(g)

    # 生成所有节点的编号列表，范围为[0, 节点总数)，这里使用第0类节点的节点数量作为总节点数。（可能是目标节点类型的数量）
    all_nodes = np.arange(features_list[0].shape[0])
    # 初始化一个形状为(节点数, 序列长度)的张量node_seq，用于存储每个节点的序列（长度为args.len_seq）。
    node_seq = torch.zeros(features_list[0].shape[0], args.len_seq).long()
    # 初始化变量n，用于遍历节点列表。
    n = 0
    # 遍历所有节点，x是当前遍历的节点编号。
    # 进行D跳上下文采样
    for x in all_nodes:
        # 初始化计数器cnt和scnt，表示序列中的位置。node_seq[n, cnt] = x表示将当前节点x放入该节点的序列中，然后增加cnt。
        cnt = 0
        scnt = 0
        node_seq[n, cnt] = x
        cnt += 1
        # 设置start为序列中第一个节点的编号，准备在接下来的步骤中查找其邻居。
        start = node_seq[n, scnt].item()
        # 当序列中的节点数小于指定的序列长度args.len_seq时，继续采样邻居节点。
        while (cnt < args.len_seq):
            # 获取节点start的所有后继节点（邻居），并转换为列表形式。
            # g.successors(start):获取图 g 中某个节点 start 的所有后继节点（即从 start 出发有直接边指向的节点）。
            # 通常，g 是某种图结构，例如 networkx 图或其他支持 .successors() 方法的图对象。
            # .numpy(): 将后继节点的集合转换为 NumPy 数组。这个操作假设 g.successors(start) 返回的数据类型支持 .numpy() 方法。
            # 如果返回的数据本身是张量或其他结构（如 torch.Tensor 或 mxnet.ndarray），可以用 .numpy() 将其转为 NumPy 数组。
            # .tolist():将 NumPy 数组转换为 Python 列表。
            sample_list = g.successors(start).numpy().tolist()
            # 确保邻居列表的长度至少为1，如果没有邻居，则设置nsampled为1。
            nsampled = max(len(sample_list), 1)
            # 随机从邻居列表sample_list中采样nsampled个节点，存储在sampled_list中。
            sampled_list = random.sample(sample_list, nsampled)
            # 遍历采样的邻居节点，将它们依次填入node_seq的第n行，同时增加cnt。如果cnt达到了args.len_seq，则退出循环。
            for i in range(nsampled):
                node_seq[n, cnt] = sampled_list[i]
                cnt += 1
                if cnt == args.len_seq:
                    break
            # 更新scnt，使start等于当前序列中的下一个节点，准备采样新的邻居。
            scnt += 1
            start = node_seq[n, scnt].item()
        # 更新n，处理下一个节点的序列生成。
        n += 1
    # 生成一个node_type列表，表示每个节点的类型。node_cnt表示每种节点类型的数量，通过嵌套的列表推导式为每个节点生成相应的类型标签。
    # for x in range(z):对于每个 (i, z) 元组，range(z) 生成从 0 到 z-1 的整数序列。这个内层循环的作用是将 i 重复 z 次。
    node_type = [i for i, z in zip(range(len(node_cnt)), node_cnt) for x in range(z)]

    # 将图g传输到指定的设备（GPU或CPU）。
    g = g.to(device)
    # 根据训练集、验证集、测试集的节点索引，分别获取对应的节点序列，并将其存储在train_seq、val_seq和test_seq中。
    train_seq = node_seq[train_idx]
    val_seq = node_seq[val_idx]
    test_seq = node_seq[test_idx]

    # 初始化两个张量micro_f1和macro_f1，每个张量的长度为args.repeat，用于存储多次实验中计算得到的微平均和宏平均F1分数。
    micro_f1 = torch.zeros(args.repeat)
    macro_f1 = torch.zeros(args.repeat)

    # 获取数据集中标签的类别数。
    num_classes = dl.labels_train['num_classes']
    # 创建一个类型嵌入矩阵type_emb，为每种节点类型创建一个单位矩阵（one-hot编码），并将其传输到设备。一个形状为 [len(node_cnt), len(node_cnt)] 的方阵
    type_emb = torch.eye(len(node_cnt)).to(device)
    # 将node_type转换为torch.Tensor类型，并传输到设备。
    node_type = torch.tensor(node_type).to(device)

    # 模型训练和评估将进行 args.repeat 次
    for i in range(args.repeat):
        # 创建 HINormer 模型实例，传入图 g、类别数 num_classes、输入维度 in_dims、隐藏层维度、层数、GNN层数、头数、dropout概率等参数。
        net = HINormer(g, num_classes, in_dims, args.hidden_dim, args.num_layers, args.num_gnns, args.num_heads,
                       args.dropout,
                       temper=args.temperature, num_type=len(node_cnt), beta=args.beta)
        # 将模型转移到指定的计算设备（例如GPU或CPU）。
        net.to(device)
        # 使用Adam优化器，设置学习率lr和权重衰减weight_decay进行参数更新。
        optimizer = torch.optim.Adam(net.parameters(), lr=args.lr, weight_decay=args.weight_decay)
        # 创建学习率调度器，使用ReduceLROnPlateau策略，当验证集损失不再降低时，调整学习率。
        # ReduceLROnPlateau 是 PyTorch 中的一种学习率调度策略，它会根据验证集的指标（如损失）在训练过程中检测模型的性能变化，当性能不再改善时，自动降低学习率。
        # optimizer: 传入训练模型时使用的优化器（例如 torch.optim.SGD, torch.optim.Adam），调度器会根据这个优化器的学习率参数进行调整。
        # mode='min': 指定调度器监控的指标是最小化的（例如，验证损失）。如果该指标不再下降，学习率将被减小。如果选择 'max'，则表示监控的指标是最大化的（例如，准确率）。
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min')

        # training loop
        # 设置训练模式
        net.train()
        # 创建早停机制early_stopping对象，使用 args.patience 设置的等待轮数。如果验证集的损失在多次迭代后不再改进，则提前停止训练。模型会保存到指定的路径。
        early_stopping = EarlyStopping(patience=args.patience, verbose=True,
                                       save_path='checkpoint/HINormer_{}_{}_{}.pt'.format(args.dataset, args.num_layers,
                                                                                          args.device))
        # 开始训练循环，每次训练epoch个周期。t_start用于记录开始训练的时间。
        for epoch in range(args.epoch):
            t_start = time.time()
            # training
            net.train()
            # 前向传播，通过模型 net 计算输入features_list和节点序列train_seq的输出 logits。type_emb 表示节点类型嵌入，node_type表示每个节点的类型，l2norm是正则化参数。
            # features_list中还是有四个矩阵，每个矩阵的行数和列数都不一定一样
            logits = net(features_list, train_seq, type_emb, node_type, args.l2norm)
            # 计算模型的输出的对数概率logp，并使用负对数似然损失函数nll_loss计算训练集上的损失train_loss。
            logp = F.log_softmax(logits, 1)
            train_loss = F.nll_loss(logp, labels[train_idx])

            # autograd
            # 清空优化器中的梯度缓存，反向传播计算梯度，并更新模型参数。
            optimizer.zero_grad()
            train_loss.backward()
            optimizer.step()
            # 记录训练时间并输出训练损失信息。
            t_end = time.time()
            # print training info
            print('Epoch {:05d} | Train_Loss: {:.4f} | Time: {:.4f}'.format(
                epoch, train_loss.item(), t_end - t_start))

            t_start = time.time()
            # validation
            # 进入验证阶段，将模型设为评估模式net.eval()，禁用梯度计算torch.no_grad()。通过验证集前向传播并计算损失val_loss。然后对模型的输出进行分类预测并转换为one-hot编码，然后评估验证集结果。
            net.eval()
            with torch.no_grad():
                logits = net(features_list, val_seq, type_emb, node_type, args.l2norm)
                logp = F.log_softmax(logits, 1)
                val_loss = F.nll_loss(logp, labels[val_idx])
                pred = logits.cpu().numpy().argmax(axis=1)
                onehot = np.eye(num_classes, dtype=np.int32)
                pred = onehot[pred]
                print(dl.evaluate_valid(pred, dl.labels_train['data'][val_idx]))

            # 更新学习率调度器，基于验证集损失val_loss调整学习率。
            scheduler.step(val_loss)
            t_end = time.time()
            # print validation info
            print('Epoch {:05d} | Val_Loss {:.4f} | Time(s) {:.4f}'.format(
                epoch, val_loss.item(), t_end - t_start))
            # early stopping
            # 检查验证集损失是否有改进，如果没有改进达到设定次数，则触发早停机制，提前结束训练。
            early_stopping(val_loss, net)
            if early_stopping.early_stop:
                print('Early stopping!')
                break

        # testing with evaluate_results_nc
        # 加载之前保存的模型参数，准备进行测试。
        net.load_state_dict(torch.load(
            'checkpoint/HINormer_{}_{}_{}.pt'.format(args.dataset, args.num_layers, args.device)))
        # 进入测试阶段，模型设为评估模式net.eval()，在测试集上进行前向传播，获取预测结果test_logits。
        net.eval()
        with torch.no_grad():
            logits = net(features_list, test_seq, type_emb, node_type, args.l2norm)
            test_logits = logits
            # 如果运行模式为1，则生成一个包含测试集预测结果的文件，用于后续评估。
            if args.mode == 1:
                pred = test_logits.cpu().numpy().argmax(axis=1)
                dl.gen_file_for_evaluate(test_idx=test_idx, label=pred, file_name=f"{args.dataset}_{i + 1}.txt")
            else:
                # 否则，计算测试集的分类结果并使用dl.evaluate_valid评估结果，输出micro-f1和macro-f1得分。
                pred = test_logits.cpu().numpy().argmax(axis=1)
                onehot = np.eye(num_classes, dtype=np.int32)
                pred = onehot[pred]
                result = dl.evaluate_valid(pred, dl.labels_test['data'][test_idx])
                print(result)
                micro_f1[i] = result['micro-f1']
                macro_f1[i] = result['macro-f1']
    # 输出多次实验的微平均micro-f1和宏平均macro-f1的均值和标准差。
    print('Micro-f1: %.4f, std: %.4f' % (micro_f1.mean().item(), micro_f1.std().item()))
    print('Macro-f1: %.4f, std: %.4f' % (macro_f1.mean().item(), macro_f1.std().item()))


if __name__ == '__main__':
    args = dblp_params()
    run_model(args)
